%-------------------------------------------------------------------------%
% This program estimates age-specific parameters describing the distrib-
% ution of wealth. 
%-------------------------------------------------------------------------%

%Set paths
cd 'U:\data'
addpath('U:\programs');
addpath('U:\output');
clc
clear

%Import raw data
% col 1 = age
% col 2 = log DC wealth 
% col 3 = log non-DC wealth
% col 4 = sample weight
data = csvread('logwealthdata.csv',1,0);

%-------------------------------------------------------------------------%
% Estimate wealth distribution parameters using MLE for each age 51-59
%-------------------------------------------------------------------------%

%table to hold estimates
J_params = zeros(9,7); %joint distribution when DC wealth > 0
A_params = zeros(9,5); %distribution of non-DC wealth when DC wealth =0 

%initial guess of joint distribution parameters
initguess = [10;10;1.5;1.5;1];

%Compute estiimates
for a = 51:59

    %Select age = a
    temp = data(data(:,1)==a,2:end);
    x = temp(:,1);
    y = temp(:,2);
    d = (x==0);
    wts = temp(:,3);
    D = [x y d wts];
    D1 = D(D(:,3)==0,:); %DC wealth > 0 
    D2 = D(D(:,3)==1,:); %DC wealth = 0

    %Estimate parameters
    f = @(p)mlmvnorm(D1,p);
    options = optimset('MaxIter',500,'Display','iter');
    [p_hat,~,~,~] = fminsearch(f,initguess,options);
    
    %Populate tables
    r = a-50;
    J_params(r,1) = a;
    J_params(r,2:6) = p_hat';
    J_params(r,7) = length(D1);
    
    A_params(r,1) = a;
    mean_hat = sum(D2(:,2).*D2(:,4))/sum(D2(:,4));
    sd_hat = sqrt(sum(D2(:,4).*(D2(:,2)-mean_hat).^2)/sum(D2(:,4)));
    A_params(r,2) = mean_hat;
    A_params(r,3) = sd_hat;
    A_params(r,4) = sum(d.*wts)/sum(wts);
    A_params(r,5) = length(D2);
end

%Export tables
J_paramtab = array2table(J_params);
J_paramtab.Properties.VariableNames(1:7)=["Age","mu1","mu2","sigma1","sigma2","sigma12","N"];
A_paramtab = array2table(A_params);
A_paramtab.Properties.VariableNames(1:5)=["Age","mu","sigma","FracDC0","N"];
writetable(J_paramtab,'U:\output\J_params.csv','WriteVariableNames',1);
writetable(A_paramtab,'U:\output\A_params.csv','WriteVariableNames',1);
    